import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from .model import mobilenet_v2

from safebooru.labels import main_chara

"""transforms"""
data_transforms = {
    'train': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]),
    'val': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]),
}
# transform = transforms.Compose(
#     [transforms.Resize((224, 224),), transforms.ToTensor()])

"""constant for classes"""

num_classes = len(main_chara)

"""batch_size"""
batch_size = 256

# print(resnet50)
"""define device"""
device = 'cuda'
"""define model"""
resnet50 = mobilenet_v2(pretrained=True, num_classes=num_classes)
resnet50 = nn.DataParallel(resnet50)
resnet50 = resnet50.to(device)
"""loss"""
criterion = nn.CrossEntropyLoss()
"""optimizer"""

optimizer_ft = optim.SGD(resnet50.parameters(), lr=0.01, momentum=0.9)
